"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：为 optuna 设计的支持多进程、不依赖数据库、支持持久化保存的高性能存储库
"""

import atexit
import copy
import os
import pickle
import threading
from typing import Container, Optional, Sequence, Any, Dict, List
from optuna import distributions
from optuna.trial import FrozenTrial, TrialState
from optuna.storages._in_memory import InMemoryStorage
from optuna.storages._base import BaseStorage
from multiprocessing.managers import SyncManager

max_trial_num = 1000000


class CompositeStorage(BaseStorage):
    """
    为 Optuna 扩展的支持读取多个存储库复合成一个虚拟存储库的类。
    """

    def __init__(self, storage_dict: Dict[str, BaseStorage]) -> None:
        super().__init__()
        self._storage_dict = storage_dict
        self._study_list = []
        for key, storage in storage_dict.items():
            item = {}
            all_studies = storage.get_all_studies()
            for study in storage.get_all_studies():
                item["study_name"] = (
                    key if (len(all_studies) == 1) else study.study_name
                )
                item["study_id"] = study._study_id
                study.study_name = item["study_name"]
                study._study_id = len(self._study_list)
                item["storage"] = storage
                item["study"] = study
                self._study_list.append(item)

    def create_new_study(self, directions, study_name) -> int:
        raise NotImplementedError

    def delete_study(self, study_id: int) -> None:
        raise NotImplementedError

    def set_study_user_attr(self, study_id: int, key: str, value: Any) -> None:
        raise NotImplementedError

    def set_study_system_attr(self, study_id: int, key: str, value) -> None:
        raise NotImplementedError

    def get_study_id_from_name(self, study_name: str) -> int:
        for n in range(0, len(self._study_list)):
            if self._study_list[n]["study_name"] == study_name:
                return n
        raise KeyError

    def get_study_name_from_id(self, study_id: int) -> str:
        if study_id >= 0 and study_id < len(self._study_list):
            return self._study_list[study_id]["study_name"]
        raise KeyError

    def get_study_directions(self, study_id: int):
        if study_id >= 0 and study_id < len(self._study_list):
            return self._study_list[study_id]["storage"].get_study_directions(
                self._study_list[study_id]["study_id"]
            )
        raise KeyError

    def get_study_user_attrs(self, study_id: int) -> Dict[str, Any]:
        if study_id >= 0 and study_id < len(self._study_list):
            return self._study_list[study_id]["storage"].get_study_user_attrs(
                self._study_list[study_id]["study_id"]
            )
        raise KeyError

    def get_study_system_attrs(self, study_id: int) -> Dict[str, Any]:
        if study_id >= 0 and study_id < len(self._study_list):
            return self._study_list[study_id]["storage"].get_study_system_attrs(
                self._study_list[study_id]["study_id"]
            )
        raise KeyError

    def get_all_studies(self):
        return reversed([study["study"] for study in self._study_list])

    def create_new_trial(
        self, study_id: int, template_trial: Optional[FrozenTrial] = None
    ) -> int:
        raise NotImplementedError

    def set_trial_param(
        self,
        trial_id: int,
        param_name: str,
        param_value_internal: float,
        distribution,
    ) -> None:
        raise NotImplementedError

    def set_trial_state_values(
        self, trial_id: int, state: TrialState, values: Optional[Sequence[float]] = None
    ) -> bool:
        raise NotImplementedError

    def set_trial_intermediate_value(
        self, trial_id: int, step: int, intermediate_value: float
    ) -> None:
        raise NotImplementedError

    def set_trial_user_attr(self, trial_id: int, key: str, value: Any) -> None:
        raise NotImplementedError

    def set_trial_system_attr(self, trial_id: int, key: str, value) -> None:
        raise NotImplementedError

    def get_trial(self, trial_id: int) -> FrozenTrial:
        study_id = trial_id // max_trial_num
        trial_id = trial_id % max_trial_num
        if study_id >= 0 and study_id < len(self._study_list):
            return self._study_list[study_id]["storage"].get_trial(trial_id)
        raise KeyError

    def get_all_trials(
        self,
        study_id: int,
        deepcopy: bool = True,
        states: Optional[Container[TrialState]] = None,
    ) -> List[FrozenTrial]:
        if study_id >= 0 and study_id < len(self._study_list):
            all_trials = self._study_list[study_id]["storage"].get_all_trials(
                self._study_list[study_id]["study_id"], deepcopy, states
            )
            for trial in all_trials:
                # 用高位部分保存 study_id
                trial._trial_id = (
                    trial._trial_id % max_trial_num + study_id * max_trial_num
                )
            return all_trials
        raise KeyError


class MultiTaskStorage(InMemoryStorage):
    """
    为 Optuna 扩展的支持多进程任务、持久化保存、无需数据库的高性能 Storage 类。
    """

    def __init__(self, save_file_name=None, autosave_interval=30) -> None:
        super().__init__()

        # __del__ 中保存可能报错，改为此处处理
        atexit.register(self.free)

        self._save_file_name = save_file_name
        self.load()

        self._autosave_interval = autosave_interval
        self._enable_autosave = False
        if (
            self._save_file_name is not None
            and self._autosave_interval is not None
            and self._autosave_interval > 0
        ):
            # 启动后台线程来保存数据
            self._stop_event = threading.Event()
            self._autosave_thread = threading.Thread(
                target=self._save_thread_func, daemon=True
            )
            self._autosave_thread.start()
            self._enable_autosave = True

    def __getstate__(self) -> Dict[Any, Any]:
        state = super().__getstate__()
        for key in [
            "_save_file_name",
            "_autosave_interval",
            "_stop_event",
            "_autosave_thread",
        ]:
            if key in state:
                del state[key]
        return state

    def load(self):
        # 从文件中加载数据
        if self._save_file_name is not None and os.path.isfile(self._save_file_name):
            with self._lock:
                self.__setstate__(pickle.load(open(self._save_file_name, "rb")))

    def save(self):
        # 保存数据到文件
        if self._save_file_name is not None:
            with self._lock:
                os.makedirs(os.path.dirname(self._save_file_name), exist_ok=True)
                pickle.dump(self.__getstate__(), open(self._save_file_name, "wb"))

    def free(self):
        # 释放对象时，退出保存数据的后台线程
        self.save()
        if self._enable_autosave:
            self._stop_event.set()
            self._autosave_thread.join()
        self._save_file_name = None

    def _save_thread_func(self):
        # 保存数据线程
        while not self._stop_event.is_set():
            self._stop_event.wait(self._autosave_interval)
            self.save()

    def get_all_trials_count(
        self,
        study_name: str,
        states: Optional[Container[TrialState]] = (
            TrialState.COMPLETE,
            TrialState.RUNNING,
        ),
    ) -> int:
        # 获取指定学习任务的 trials 数量
        study_id = super().get_study_id_from_name(study_name)
        return len(super().get_all_trials(study_id, False, states))

    def get_diff_trials(
        self,
        study_id: int,
        cached_max_number: int,
        unfinished_ids: List[int],
        states: Optional[Container[TrialState]] = None,
    ) -> Dict[int, FrozenTrial]:
        # 获取指定 study_id 的所有未完成的 trials 和新增的 trials，实现增量更新提升性能
        with self._lock:
            self._check_study_id(study_id)

            # 取得所有需要更新的 trial id 列表
            trials = self._studies[study_id].trials
            diff_list = unfinished_ids + [
                id for id in range(cached_max_number + 1, len(trials))
            ]
            # 返回指定状态的 trial 列表
            diff_dict = {
                id: copy.deepcopy(trials[id])
                for id in diff_list
                if states is None or trials[id].state in states
            }

        return diff_dict

    def set_trial_param_and_state(
        self,
        trial_id: int,
        param_dict: Dict[str, Any],
        user_attrs: Dict[str, Any],
        state: TrialState,
        values: Optional[Sequence[float]] = None,
    ) -> bool:
        with self._lock:
            trial = self._get_trial(trial_id)

            self.check_trial_is_updatable(trial_id, trial.state)

            study_id = self._trial_id_to_study_id_and_number[trial_id][0]
            for param_name, (param_value_internal, distribution) in param_dict.items():
                # Check param distribution compatibility with previous trial(s).
                if param_name in self._studies[study_id].param_distribution:
                    distributions.check_distribution_compatibility(
                        self._studies[study_id].param_distribution[param_name],
                        distribution,
                    )

                # Set param distribution.
                self._studies[study_id].param_distribution[param_name] = distribution

            # Set param.
            trial = copy.copy(trial)
            trial.params = copy.copy(trial.params)
            for param_name, (param_value_internal, distribution) in param_dict.items():
                trial.params[param_name] = distribution.to_external_repr(
                    param_value_internal
                )
                trial.distributions = copy.copy(trial.distributions)
                trial.distributions[param_name] = distribution

            # Set user_attrs
            if user_attrs:
                trial.user_attrs = copy.copy(trial.user_attrs)
                trial.user_attrs.update(user_attrs)

            self._set_trial(trial_id, trial)

            return self.set_trial_state_values(trial_id, state, values)


#######################################################
# 优化多进程通讯，减少数据交换次数和数量，提升整体性能的函数
#######################################################


def __client_get_all_trials(
    self,
    study_id: int,
    deepcopy: bool = True,
    states: Optional[Container[TrialState]] = None,
) -> List[FrozenTrial]:
    # 用于替换任务进程中代理 Storage 对象的 get_all_trials() 方法，实现增量更新 trials
    # 初始化对象，取得缓存字典
    if not hasattr(self, "__all_trials_local_cache"):
        setattr(self, "__all_trials_local_cache", {})
    if not study_id in self.__all_trials_local_cache.keys():
        self.__all_trials_local_cache[study_id] = []
    trials = self.__all_trials_local_cache[study_id]

    # 从服务器取得差异 trial 列表
    cached_max_number = len(trials) - 1
    unfinished_ids = [
        id
        for id in range(len(trials))
        if trials[id] is None or not trials[id].state.is_finished()
    ]
    diff_dict = self.get_diff_trials(
        study_id, cached_max_number, unfinished_ids, states
    )

    # 如果有新 trial，先更新缓存列表长度
    if len(diff_dict) > 0:
        new_max_number = max(list(diff_dict.keys()))
        if new_max_number > cached_max_number:
            trials.extend([None] * (new_max_number - cached_max_number))

    # 用新读到的 trial 更新缓存列表
    for id, trial in diff_dict.items():
        trials[id] = trial
    self.__all_trials_local_cache[study_id] = trials

    # 取出符合状态的 trial 列表
    if states is not None:
        trials = [t for t in trials if (t is not None) and (t.state in states)]

    if deepcopy:
        trials = copy.deepcopy(trials)

    return trials


def __client_set_trial_param(
    self,
    trial_id: int,
    param_name: str,
    param_value_internal: float,
    distribution: distributions.BaseDistribution,
) -> None:
    # 用于替换任务进程中代理 Storage 对象的 set_trial_param() 方法，实现批量提交
    # 初始化对象，取得缓存字典
    if not hasattr(self, "__trial_param_local_cache"):
        setattr(self, "__trial_param_local_cache", {})
    if not trial_id in self.__trial_param_local_cache.keys():
        self.__trial_param_local_cache[trial_id] = {}
    # 测试参数只有测试结束后才有用，在任务进程中先临时缓存起来，等任务结束后再批量提交
    param_dict = self.__trial_param_local_cache[trial_id]
    param_dict[param_name] = (param_value_internal, copy.deepcopy(distribution))
    self.__trial_param_local_cache[trial_id] = param_dict


def __client_set_trial_user_attr(
    self,
    trial_id: int,
    key: str,
    value: Any,
) -> None:
    # 用于替换任务进程中代理 Storage 对象的 set_trial_user_attr() 方法，实现批量提交
    # 初始化对象，取得缓存字典
    if not hasattr(self, "__trial_user_attr_local_cache"):
        setattr(self, "__trial_user_attr_local_cache", {})
    if not trial_id in self.__trial_user_attr_local_cache.keys():
        self.__trial_user_attr_local_cache[trial_id] = {}
    # 测试参数只有测试结束后才有用，在任务进程中先临时缓存起来，等任务结束后再批量提交
    user_attr_dict = self.__trial_user_attr_local_cache[trial_id]
    user_attr_dict[key] = value
    self.__trial_user_attr_local_cache[trial_id] = user_attr_dict


def __client_set_trial_state_values(
    self, trial_id: int, state: TrialState, values: Optional[Sequence[float]] = None
) -> bool:
    # 用于替换任务进程中代理 Storage 对象的 set_trial_state_values() 方法，实现批量提交
    if (
        state.is_finished()
        and hasattr(self, "__trial_param_local_cache")
        and (trial_id in self.__trial_param_local_cache.keys())
    ):
        # 读取用户参数
        if hasattr(self, "__trial_user_attr_local_cache") and (
            trial_id in self.__trial_user_attr_local_cache.keys()
        ):
            user_attrs = self.__trial_user_attr_local_cache.pop(trial_id)
        else:
            user_attrs = {}

        # 只有已经缓存了参数并且执行完成的测试才使用批量提交，使用 pop 从字典中弹出数据，避免内存泄漏
        return self.set_trial_param_and_state(
            trial_id, self.__trial_param_local_cache.pop(trial_id), user_attrs, state, values
        )
    else:
        # 其它情况下调用已保存的 set_trial_state_values() 方法
        return self.__saved_set_trial_state_values(trial_id, state, values)


def __hook_obj_method(obj, method_name, new_method):
    # 将对象的方法替换为新方法，并保存原来的方法
    saved_name = f"__saved_{method_name}"
    if not hasattr(obj, saved_name):
        setattr(obj, saved_name, getattr(obj, method_name))
        setattr(obj, method_name, new_method.__get__(obj))


def get_fast_client_storage(storage, fast_get_trials=True, fast_set_trial_param=True):
    # 替换任务进程中 storage 代理对象中的方法，提升性能
    if fast_get_trials:
        __hook_obj_method(storage, "get_all_trials", __client_get_all_trials)
    if fast_set_trial_param:
        __hook_obj_method(storage, "set_trial_param", __client_set_trial_param)
        __hook_obj_method(storage, "set_trial_user_attr", __client_set_trial_user_attr)
        __hook_obj_method(
            storage, "set_trial_state_values", __client_set_trial_state_values
        )
    return storage


# 注册成支持多进程共享的对象
SyncManager.register("InMemoryStorage", InMemoryStorage)
SyncManager.register("MultiTaskStorage", MultiTaskStorage)
